# -*- coding: utf-8 -*-

import numpy as np
import random
import rospy
import sys
import importlib.util
import os
import copy

from pvesti.msg import UAVPose, UAVPosVel
from geometry_msgs.msg import Pose, Point, Twist, Vector3
from geometry_msgs.msg import PoseStamped
from nav_msgs.msg import Odometry



from kalmanPredict import KalmanFilterPredictor

# 检测到的无人机个数
other_iris_names = ['iris_1']
# other_iris_names = []
OTHER_UAV_NUM = len(other_iris_names)

# 当前无人机卡尔曼滤波器
current_uav_kalman_predictor = None 
other_uav_kalman_predictors = {}


lio_pos = Point()        
lio_rot = np.eye(3)      
lio_timestamp = rospy.Time(0)


lio_other_poses = [Point()] * OTHER_UAV_NUM
lio_other_rots = [np.eye(3)] * OTHER_UAV_NUM
lio_other_timestamps = [rospy.Time(0)] * OTHER_UAV_NUM



def quaternion_to_rotation_matrix(q):

    [q_x, q_y, q_z, q_w] = q

    # 计算旋转矩阵
    R = np.array([
        [1 - 2 * (q_y**2 + q_z**2), 2 * (q_x * q_y - q_z * q_w), 2 * (q_x * q_z + q_y * q_w)],
        [2 * (q_x * q_y + q_z * q_w), 1 - 2 * (q_x**2 + q_z**2), 2 * (q_y * q_z - q_x * q_w)],
        [2 * (q_x * q_z - q_y * q_w), 2 * (q_y * q_z + q_x * q_w), 1 - 2 * (q_x**2 + q_y**2)]
    ])
    
    return R



# UAV 绝对位置信息回调函数
def lio_self_callback(data):

    # 整型类型要修改一定需要global字段，列表不需要
    global lio_timestamp, lio_pos, lio_rot
    
    lio_timestamp = data.header.stamp
    lio_pos = data.pose.pose.position
    # lio_pos.x = lio_pos.x + BIAS[0]
    # lio_pos.y = lio_pos.y + BIAS[1]
    # lio_pos.z = lio_pos.z + BIAS[2]
    
    lio_quat = [data.pose.pose.orientation.x, data.pose.pose.orientation.y, 
                 data.pose.pose.orientation.z, data.pose.pose.orientation.w]
    lio_rot = quaternion_to_rotation_matrix(lio_quat)
    

def lio_others_callback(data, i):
    # global lio_pos_other, lio_rot_other

    lio_other_timestamps[i] = data.header.stamp
    lio_other_poses[i] = data.pose.position
    
    lio_quat = [data.pose.orientation.x, data.pose.orientation.y, 
                 data.pose.orientation.z, data.pose.orientation.w]
    lio_other_rots[i] = quaternion_to_rotation_matrix(lio_quat)
    
    

# 主函数，用于 ROS 节点，接收 UAV 位置信息并发布位置信息和速度
def main():
    
    iris_name = sys.argv[1]
    
    
    rospy.init_node(iris_name + 'pos_vel_transfer')
    
    # 订阅当前无人机的话题
    rospy.Subscriber('/Odometry', Odometry,lio_self_callback, queue_size=1)
    # 订阅roslibpy转化的其他无人机的话题
    for index, other_iris_name in enumerate(other_iris_names):
        topic_name = '/transfer/' + other_iris_name + '/mavros/vision_pose/pose'
        rospy.loginfo("topic name': {}".format(topic_name))
        rospy.Subscriber(topic_name, PoseStamped, lio_others_callback, index, queue_size=1)
    # rospy.Subscriber(vehicle_type+'_'+str(i)+'/vins_estimator/odometry', Odometry,vins_callback, queue_size=1)#从vins中得到中心无人机的绝对位置信息
    
    # publish all uav pose including itself
    pub = rospy.Publisher('/' + iris_name + '/uav_pos_vel', UAVPosVel, queue_size=10)#给处理避障的节点发布无人机数量、位置、速度信息uav_count、geometry_msgs/Pose、geometry_msgs/Twist
    

    global lio_pos, lio_timestamp
    
    first_run = True
    ROS_RATE = 10
    rate = rospy.Rate(ROS_RATE)
    
    last_lio_pos = [0,0,0]
    # 卡尔曼的估计频率就是dt，但是是预测速度+
    while not rospy.is_shutdown():
        
        dt = 1.0/ROS_RATE
        
        # 当前uav卡尔曼滤波器-预测
        if first_run:
            # 初始化当前uav的卡尔曼滤波器 
            current_uav_kalman_predictor = KalmanFilterPredictor(dt)
            current_uav_kalman_predictor.initialize_state([lio_pos.x, lio_pos.y, lio_pos.z])
            # 初始化其他uavs的卡尔曼滤波器 
            for index, other_iris_name in enumerate(other_iris_names):
                other_uav_kalman_predictors[index] = KalmanFilterPredictor(dt)
                other_uav_kalman_predictors[index].initialize_state([lio_other_poses[index].x, lio_other_poses[index].y, lio_other_poses[index].z])

            
            first_run = False
            
            
        current_position = [lio_pos.x, lio_pos.y, lio_pos.z]
        current_uav_kalman_predictor.predict(current_position, dt)
        
        for index in range(OTHER_UAV_NUM):
            other_position = [lio_other_poses[index].x, lio_other_poses[index].y, lio_other_poses[index].z]
            other_uav_kalman_predictors[index].predict(other_position, dt)
            rospy.loginfo("{} other_lio': {}".format(other_iris_names[index], other_position))
        
        msg = UAVPosVel()
        msg.header.stamp = lio_timestamp
        msg.uav_count = OTHER_UAV_NUM + 1
        
        # 增加当前无人机的位置和速度
        pose_cur = Pose()
        pose_cur.position = Point(*current_uav_kalman_predictor.filtered_positions[-1])
        twist_cur = Twist()
        twist_cur.linear = Vector3(*current_uav_kalman_predictor.filtered_velocities[-1])
        msg.uav_positions.append(pose_cur)
        msg.uav_velocities.append(twist_cur)
        
        pose_cur_info = [pose_cur.position.x, pose_cur.position.y, pose_cur.position.z]
        vel_cur_info = [twist_cur.linear.x, twist_cur.linear.y, twist_cur.linear.z]
        lio_vel = ((np.array(current_position) - np.array(last_lio_pos))/dt).tolist()
        rospy.loginfo("lio Odometry': {}".format(current_position))
        # rospy.loginfo("lio Velocity': {}".format(lio_vel))
        rospy.loginfo("current UAV position': {}".format(pose_cur_info))
        rospy.loginfo("current UAV velocity': {}".format(vel_cur_info))
        
        
        # 增加roslibpy获得的其他无人机的位置和速度
        for index in range(OTHER_UAV_NUM):
            pose_other = Pose()
            pose_other.position = Point(*other_uav_kalman_predictors[index].filtered_positions[-1])
            twist_other = Twist()
            twist_other.linear = Vector3(*other_uav_kalman_predictors[index].filtered_velocities[-1])
            msg.uav_positions.append(pose_other)
            msg.uav_velocities.append(twist_other)
            
            pose_other_info = [pose_other.position.x, pose_other.position.y, pose_other.position.z]
            vel_other_info = [twist_other.linear.x, twist_other.linear.y, twist_other.linear.z]
            rospy.loginfo("{} position': {}".format(other_iris_names[index], pose_other_info))
            rospy.loginfo("{} velocity': {}".format(other_iris_names[index], vel_other_info))
        

        last_lio_pos = current_position
        # 发布消息    
        pub.publish(msg)
        

        rate.sleep()

if __name__ == "__main__":
    main()